#include <Eigen/Core>
#include <ceres/ceres.h>
#include <ceres/rotation.h>

#include <vector>

#include <autodiff/reverse.hpp>
#include <autodiff/reverse/eigen.hpp>

using namespace Eigen;

// 相平面坐标
std::vector<Vector2d> point2ds = {Vector2d(-86.15, -68.99), Vector2d(-53.40, 82.21), Vector2d(-14.78, -76.63),
                                  Vector2d(10.46, 64.43)};
// 控制点坐标
std::vector<Vector3d> point3ds = {Vector3d(36589.41, 25273.32, 2195.17), Vector3d(37631.08, 31324.51, 728.69),
                                  Vector3d(39100.97, 24934.98, 2386.50), Vector3d(40426.54, 30319.81, 757.31)};
// 内方位元素
double f = 153.24;
double x_0 = 0.0;
double y_0 = 0.0;

using namespace autodiff;

Matrix3var rodrigues(const Vector3var &angle_axis) {
    static const double kOne = 1.0;
    const var theta2 = angle_axis.squaredNorm();
    Matrix3var R;
    if (theta2 > var(std::numeric_limits<double>::epsilon())) {
        // We want to be careful to only evaluate the square root if the
        // norm of the angle_axis vector is greater than zero. Otherwise
        // we get a division by zero.
        const var theta = sqrt(theta2);
        const var wx = angle_axis[0] / theta;
        const var wy = angle_axis[1] / theta;
        const var wz = angle_axis[2] / theta;

        const var costheta = cos(theta);
        const var sintheta = sin(theta);

        R(0, 0) = costheta + wx * wx * (kOne - costheta);
        R(1, 0) = wz * sintheta + wx * wy * (kOne - costheta);
        R(2, 0) = -wy * sintheta + wx * wz * (kOne - costheta);
        R(0, 1) = wx * wy * (kOne - costheta) - wz * sintheta;
        R(1, 1) = costheta + wy * wy * (kOne - costheta);
        R(2, 1) = wx * sintheta + wy * wz * (kOne - costheta);
        R(0, 2) = wy * sintheta + wx * wz * (kOne - costheta);
        R(1, 2) = -wx * sintheta + wy * wz * (kOne - costheta);
        R(2, 2) = costheta + wz * wz * (kOne - costheta);
    } else {
        // Near zero, we switch to using the first order Taylor expansion.
        R(0, 0) = kOne;
        R(1, 0) = angle_axis[2];
        R(2, 0) = -angle_axis[1];
        R(0, 1) = -angle_axis[2];
        R(1, 1) = kOne;
        R(2, 1) = angle_axis[0];
        R(0, 2) = angle_axis[1];
        R(1, 2) = -angle_axis[0];
        R(2, 2) = kOne;
    }
    return R;
}
Matrix3d cal_Matrix_R(const double phi, const double w, const double k) {
    Matrix3d R;
    R(0, 0) = cos(phi) * cos(k) - sin(phi) * sin(w) * sin(k);
    R(0, 1) = -cos(phi) * sin(k) - sin(phi) * sin(w) * cos(k);
    R(0, 2) = -sin(phi) * cos(w);
    R(1, 0) = cos(w) * sin(k);
    R(1, 1) = cos(w) * cos(k);
    R(1, 2) = -sin(w);
    R(2, 0) = sin(phi) * cos(k) + cos(phi) * sin(w) * sin(k);
    R(2, 1) = -sin(phi) * sin(k) + cos(phi) * sin(w) * cos(k);
    R(2, 2) = cos(phi) * cos(w);
    return R;
}
MatrixXd cal_Matrix_A(Matrix3d &R, Vector2d &P, double Z, double phi, double w, double k) {
    MatrixXd A(2, 6);
    A(0, 0) = (R(0, 0) * f + R(0, 2) * P.x()) / Z;                                                     // a11
    A(0, 1) = (R(1, 0) * f + R(1, 2) * P.x()) / Z;                                                     // a12
    A(0, 2) = (R(2, 0) * f + R(2, 2) * P.x()) / Z;                                                     // a13
    A(0, 3) = P.y() * sin(w) - (P.x() * (P.x() * cos(k) - P.y() * sin(k)) / f + f * cos(k)) * cos(w);  // a14
    A(0, 4) = -f * sin(k) - P.x() * (P.x() * sin(k) + P.y() * cos(k)) / f;                             // a15
    A(0, 5) = P.y();                                                                                   // a16
    A(1, 0) = (R(0, 1) * f + R(0, 2) * P.y()) / Z;                                                     // a21
    A(1, 1) = (R(1, 1) * f + R(1, 2) * P.y()) / Z;                                                     // a22
    A(1, 2) = (R(2, 1) * f + R(2, 2) * P.y()) / Z;                                                     // a23
    A(1, 3) = -P.x() * sin(w) - (P.y() * (P.x() * cos(k) - P.y() * sin(k)) / f - f * sin(k)) * cos(w); // a24
    A(1, 4) = -f * cos(k) - (P.y() * (P.x() * sin(k) + P.y() * cos(k))) / f;                           // a25
    A(1, 5) = -P.x();                                                                                  // a26
    return A;
}

inline Vector2d projectp(const Vector3d &point, const Vector3d &Xs, const Matrix3d &R) {
    Vector2d pixel;

    Vector3d X = point - Xs;
    X = R.transpose() * X;

    pixel.x() = -f * X(0) / X(2) + x_0;
    pixel.y() = -f * X(1) / X(2) + y_0;

    return pixel;
}

std::tuple<Vector3d, Matrix3d, double, Vector3d> question4() {
    Vector3d Xs = Vector3d::Zero();
    Matrix3d R = Matrix3d::Identity();
    double sigma_0 = 0.0;
    Vector3d m_Xs = Vector3d::Zero();

    // TODO 实现后方交会

    // 1 计算外方位线元素初值
    for (const auto &point : point3ds) {
        Xs += point;
    }
    Xs = Xs / point3ds.size();
    Xs.z() = 8000.0;

    // 2 角元素初值就为单位阵
    double phi = 0, w = 0, k = 0;
    MatrixXd A(point2ds.size() * 2, 6);
    VectorXd l(point2ds.size() * 2);
    VectorXd dx(6);
    Vector2d xt;

    // 3 迭代优化求解
    // 误差方程 V=A*dx-l
    for (int iter = 0; iter < 125; iter++) {
        R = cal_Matrix_R(phi, w, k);
        for (int i = 0; i < point2ds.size(); i++) {
            xt = projectp(point3ds[i], Xs, R);
            double Z = R(0, 2) * (point3ds[i].x() - Xs.x()) + R(1, 2) * (point3ds[i].y() - Xs.y()) +
                       R(2, 2) * (point3ds[i].z() - Xs.z()); //便于计算矩阵A
            A.block(i * 2, 0, 2, 6) << cal_Matrix_A(R, point2ds[i], Z, phi, w, k);
            l.block(i * 2, 0, 2, 1) << point2ds[i] - xt;
        }
        // solve normal equation
        dx << (A.transpose() * A).ldlt().solve(A.transpose() * l);
        // std::cout << "A=" << A << std::endl;
        // std::cout << "l=" << l << std::endl;
        // std::cout << "dx=" << dx << std::endl;
        // std::cout << "-----------" << std::endl;
        Xs.x() += dx[0];
        Xs.y() += dx[1];
        Xs.z() += dx[2];
        phi += dx[3];
        w += dx[4];
        k += dx[5];
    }
    VectorXd v = A * dx - l;

    sigma_0 = v.norm() / std::sqrt((2 * point2ds.size() - 6));
    MatrixXd cov_Xs = (A.transpose() * A).inverse();
    m_Xs.x() = sigma_0 * std::sqrt(cov_Xs(0, 0));
    m_Xs.y() = sigma_0 * std::sqrt(cov_Xs(1, 1));
    m_Xs.z() = sigma_0 * std::sqrt(cov_Xs(2, 2));
    return std::make_tuple(Xs, R, sigma_0, m_Xs);
}

struct ProjectionResidual {
    ProjectionResidual(const Vector2d &point2d, const Vector3d &point3d) : point2d(point2d), point3d(point3d) {}

    template <typename T> bool operator()(const T *Xs, const T *r, T *residual) const {
        T R[9];
        ceres::AngleAxisToRotationMatrix(r, ceres::RowMajorAdapter3x3(R));

        T X = point3d.x() - Xs[0];
        T Y = point3d.y() - Xs[1];
        T Z = point3d.z() - Xs[2];

        T x = -f * (R[0] * X + R[3] * Y + R[6] * Z) / (R[2] * X + R[5] * Y + R[8] * Z) + x_0;
        T y = -f * (R[1] * X + R[4] * Y + R[7] * Z) / (R[2] * X + R[5] * Y + R[8] * Z) + y_0;

        residual[0] = point2d.x() - x;
        residual[1] = point2d.y() - y;

        return true;
    }

private:
    Vector2d point2d;
    Vector3d point3d;
};

std::tuple<Vector3d, Matrix3d, double, Vector3d> golden4() {
    Vector3d Xs = Vector3d::Zero();
    Matrix3d R = Matrix3d::Identity();
    double sigma_0 = 0.0;
    Vector3d m_Xs = Vector3d::Zero();

    // 计算初值
    for (const auto &point : point3ds) {
        Xs += point;
    }
    Xs = Xs / point3ds.size();
    Xs.z() = 8000.0;

    // 角元素初值就是单位阵
    auto angle_axis = Eigen::AngleAxisd(R);
    Vector3d r = angle_axis.angle() * angle_axis.axis();

    ceres::Problem problem;

    std::vector<ceres::ResidualBlockId> blocks;

    for (int i = 0; i < point3ds.size(); ++i) {
        ceres::CostFunction *cost = new ceres::AutoDiffCostFunction<ProjectionResidual, 2, 3, 3>(
            new ProjectionResidual(point2ds[i], point3ds[i]));
        ceres::ResidualBlockId block = problem.AddResidualBlock(cost, nullptr, Xs.data(), r.data());
        blocks.push_back(block);
    }

    ceres::Solver::Options op;
    op.max_num_iterations = 100;
    op.linear_solver_type = ceres::DENSE_QR;

    ceres::Solver::Summary summary;
    ceres::Solve(op, &problem, &summary);

    // 计算单位权中误差
    {
        ceres::Problem::EvaluateOptions op_eval;
        op_eval.residual_blocks = blocks;
        double cost;
        problem.Evaluate(op_eval, &cost, nullptr, nullptr, nullptr);
        cost *= 2.0; // ceres 记录的是包含 1/2 的
        sigma_0 = std::sqrt(cost / (2 * point2ds.size() - 6));
    }

    // 计算协方差矩阵
    {
        Matrix3d cov_Xs;
        std::vector<std::pair<const double *, const double *>> covariance_blocks;
        covariance_blocks.emplace_back(Xs.data(), Xs.data());
        ceres::Covariance::Options options;
        ceres::Covariance covariance(options);
        covariance.Compute(covariance_blocks, &problem);
        covariance.GetCovarianceBlock(Xs.data(), Xs.data(), cov_Xs.data());

        m_Xs.x() = sigma_0 * std::sqrt(cov_Xs(0, 0));
        m_Xs.y() = sigma_0 * std::sqrt(cov_Xs(1, 1));
        m_Xs.z() = sigma_0 * std::sqrt(cov_Xs(2, 2));
    }

    // 计算旋转矩阵
    {
        angle_axis.angle() = r.norm();
        angle_axis.axis() = r.normalized();

        R = angle_axis.toRotationMatrix();
    }
    return std::make_tuple(Xs, R, sigma_0, m_Xs);
}
